"""
PySCeS - Python Simulator for Cellular Systems
           (http://pysces.sourceforge.net)
Copyright (C) B.G. Olivier, J.M. Rohwer, J.-H.S. Hofmeyr
Stellenbosch, 2004-2006.
Triple-J Group for Molecular Cell Physiology
Stellenbosch University, South Africa

PySCeS is Open Source Software distributed under
the GNU GENERAL PUBLIC LICENSE (see docs/GPL)

"""

__doc__ = '''SBML reading/writing module'''
__version__ = 'development'
__lastauthor__ = '$Author: c.gillespie@ncl.ac.uk $'
__revision__ = '$Rev: 54 $'

import os, sys
from time import strftime
from getpass import getuser

if sys.platform == 'win32':
    try:
        import pysces.libsbml.libsbmlinit as libsbml
    except ImportError, e:
        print 'Windows sbml load error', e
else:
    try:
        import libsbml
    except ImportError, e:
        print 'Posix sbml load error', e



class SBML2psc(object):
    
    mode_number_format = '%2.4e'
    
    def __init__(self, sbml_string=None, sbml_file=None, dir_in=None):
        r = libsbml.SBMLReader()

        if sbml_string:
            self.d = r.readSBMLFromString(sbml_string)
        else:
            if dir_in == None:
                dir_in = os.getcwd()
            file_path = os.path.join(dir_in, sbml_file)
            assert os.path.exists(file_path), "Invalid path"
            self.d = r.readSBML(file_path)
            

        self.d.setLevel(2)
        self.d.setVersion(1)
        self.m = self.d.getModel()
        self.psc_string = ''
        self.model_filename = file_path

    def writePSCToFile(self, file_name = None, dir_out=None):
        '''
            writePSCToFile()

            Writes the PSC model to a file
            Arguments:
            =========
            file_name: output file
            dir_out output directory
        '''

        if not dir_out:
            dir_out = os.getcwd()
        
        if file_name:
            psc_filename = os.path.join(dir_out, file_name)
        else:
            psc_filename = os.path.join(dir_out, 'tmp.psc')
            
        output_file = open(psc_filename,'w')
        output_file.write(self.__getPSCString())
        output_file.close()
        return True

    def writePSCToString(self):
        '''
            writePSCToFile()

            Writes the PSC model to a file
        '''

        self.__getPSCString()
        return self.psc_string
        
                      

        
    def __getPSCString(self):
        """
        __getPSCString(filename,dir=None,dirOut=None)

        Convert an SBML file into a PySCeS string
        """        
        #Create list of reaction ids      
        reaction_ids = [i.getId() for i in self.m.getListOfReactions()]
                
        #Create a list list of species
        init_fixed = []
        init_var = []
        for i in self.m.getListOfSpecies():
            parName = i.getId()
        # if a species is a BoundaryCondition or constant it becomes fixed - brett 20050111
            if i.getBoundaryCondition() or i.getConstant():
                if i.getConstant() and not i.getBoundaryCondition():
                    print parName, ' is set as constant, assuming: BoundaryCondition = True'
                init_fixed.append((parName, i.getInitialConcentration()))
            else:
                init_var.append((parName, i.getInitialConcentration()))

        NetworkDict = dict([(i, dict.fromkeys(['Params',
                              'RateEq',
                              'Reagents',
                              'Type'])) for i in reaction_ids])
        
        parameters = []
        init_par = []
        for i in self.m.getListOfReactions():
            rDict = NetworkDict[i.getId()]
            j = i.getKineticLaw()
            par = []
            try:
                for k in j.getListOfParameters():
                    par.append(k.getId())
                    init_par.append((k.getId(), k.getValue()))
                    parameters.append(k.getId())
                rDict['Params'] = par
                rDict['RateEq'] = j.getFormula()
            except Exception, err:
                rDict['Params'] = []
                rDict['RateEq'] = ''
                print err,'XXX'

            # this is to eliminate zero stoichiometries {0}xyz
            Substrates = [(k.getSpecies(), -k.getStoichiometry()) for k in i.getListOfReactants()]
            bad_list = [sub for sub in Substrates if sub[1] == 0]
            for bad in bad_list:
                Substrates.pop(Substrates.index(bad))
            if not Substrates:
                Substrates.append(('$pool', -1.0))
                
            
            Products = [(k.getSpecies(), k.getStoichiometry()) for k in i.getListOfProducts()]
            bad_list = [prob for prod in Products if prod[1] == 0]
            for bad in bad_list:
                Products.pop(Products.index(bad))
            if not Products:
                Products.append(('$pool', 1.0))


            rDict['Reagents'] = dict(Substrates + Products)
            if i.getReversible():
                t = 'Rever'
            else:
                t = 'Irrev'
            rDict['Type'] = t
            NetworkDict[i.getId()].update(rDict)

                
        # Add extra model parameters not defined in reactions (apparently)
        #XXX THIS IS WRONG  - local parameters may not be unique.
        #XXX - You may have to change the parameter name in the kinetic law.
        #XXX - SHould we use the pysbml ModelHandler class for all this conversion?
        for pa in self.m.getListOfParameters():
            if pa.getId() not in parameters:
                init_par.append((pa.getId(), pa.getValue()))

        
        self.__getPSCHeader()
        self.__getPSCFixedSpeciesList(init_fixed)
        self.__getPSCRateEquations(NetworkDict, number_format='%2.3f')
        self.__getPSCSpecies(init_var)
        self.__getPSCFixedSpecies(init_fixed)
        self.__getPSCParameters(init_par)
        self.__getPSCCompartments()
        self.__getPSCRules()
        self.__getPSCNotes()
        
        return self.psc_string 
                
       
    def __getPSCCompartments(self):
        compartmentList = []
        for comp in self.m.getListOfCompartments():
            compartmentList.append((comp.getId(), comp.getVolume()))

        if len(compartmentList) > 1:
            print '\nINFO: PySCeS models are assumed to have a single compartment'

        if len(compartmentList) > 0:
            
            self.psc_string += '\n## Initialise compartment volumes'
            for comp in compartmentList:
                self.psc_string += '\n' + comp[0] + ' = ' + str(comp[1])
                ##  parameters.append(x[0])
                ##  init_par.append(x)
            self.psc_string += '\n'

    def __getPSCRules(self):
        # Add assignment rules as forcing functions - brett 20050908
        pscRules = [(rule.getVariable(), rule.getFormula()) for rule in self.m.getListOfRules()]
        if pscRules:
            self.psc_string += '\n## Assignment rules translated to forcing functions\n'
        for rule in pscRules:
            rule0 = 'self.' + rule[0]
            rule1l = rule[1].split()
            for word in rule1l:
                if word.isalnum():
                    if word not in range(1, 10):
                        word = 'self.' + word
            self.psc_string += '!F '+ rule0 + ' = '
            for word in rule1l:
                self.psc_string += word + ' '
            self.psc_string += '\n'
        self.psc_string += '\n'
        
    def __getPSCNotes(self):
        #Check for model Notes
        if len(self.m.getNotes()) > 0:
            self.psc_string += '\n## Model notes' + self.m.getNotes().replace('\n', '\n# ') + '\n\n'
       

    def __getPSCHeader(self):
        """
        __getPSCHeader()

        Write a PSC file header to an open file object
        """        
        try:
            UseR = getuser()
        except:
            UseR = ''

        self.psc_string = ''
        #self.psc_string +=  '############################################################\n'
        self.psc_string += '# PySCeS (' + __version__ + ') model input file\n'
        self.psc_string +=  '# PySCeS can be found at http://pysces.sourceforge.net/\n'
        #self.psc_string += '###########################################################\n\n'
        self.psc_string += '# Original input file: ' + self.model_filename.split('\\')[-1][:-4] + '\n'
        self.psc_string += '# This file generated: ' + strftime("%a, %d %b %Y %H:%M:%S") + ' by '+UseR+'\n\n'
        return True


    def __getPSCSpecies(self, species):
        """
        __getPSCSpecies(species)
        Write out model species initiailisations to file

        Arguments:
        =========
        species: a list of (species.value) pairs
        """

        if not species:
            return True
        self.psc_string += '\n## Variable species initial values\n'
      
        for sp in species:
            self.psc_string += '%s = %s \n' % (sp[0], self.mode_number_format % sp[1])
            
        return True

    def __getPSCFixedSpeciesList(self, fixed_species):
        """
        __getPSCFixedSpeciesList(fixed_species)

        Write fixed species declaration to a PSC file

        Arguments:
        =========
        fixed_species: a list of (species,value) pairs

        """        
        self.psc_string += '## Fixed species\n'
        if fixed_species:
            self.psc_string += 'FIX: '
            for x in fixed_species:
                self.psc_string += x[0] + ' '
        else:     
            self.psc_string += '#  <none>'
        self.psc_string += '\n\n'


    def __getPSCFixedSpecies(self, fixed_species):
        """
        __getPSCFixedSpecies(fixed_species)

        Write fixed species initialisations to a PSC file

        Arguments:
        =========
        fixed_species: a list of (species,value) pairs

        """
        out_list = ['\n## Fixed species\n']
        for x in fixed_species:
            out_list.append(x[0] + ' = ' +  self.mode_number_format % x[1] + '\n')
        for x in out_list:
            self.psc_string += x
        self.psc_string += '\n'


    def __getPSCParameters(self, parameters):
        """
        __getPSCParameters(File,parameters)

        Write mode parameter initialisations to a PSC file

        Arguments:
        =========
        parameters: a list of (parameter,value) pairs

        """        
        out_list = ['\n## Parameters\n']
        for x in parameters:
            out_list.append(x[0] + ' = ' +  self.mode_number_format % x[1] + '\n')
        for x in out_list:
            self.psc_string += x
        self.psc_string += '\n'


    def __getPSCRateEquations(self, NetworkDict, number_format='%2.3f'):
        """
        __getPSCRateEquations(NetworkDict,number_format='%2.3f')

        Write model rate equations to a PSC file

        Arguments:
        =========
        NetworkDict: a PySCeS network dictionary
        number_format [default='%2.3f']: number formatting to use in rate laws
        """
        out_list = ['\n## Reaction stoichiometry and rate equations\n']
        for key in NetworkDict:
            out_list.append(key + ':\n')
            
            reag_l, reag_r = [], []
            reagents = NetworkDict[key]['Reagents']
            for reagent in reagents:
                if reagents[reagent] > 0:
                    if reagents[reagent] == 1.0:
                        reag_r.append(reagent.replace('self.',''))
                    else:
                        reag_r.append('{' + number_format % abs(reagents[reagent]) + '}' + reagent.replace('self.',''))
                elif reagents[reagent] < 0:
                    if reagents[reagent] == -1.0:
                        reag_l.append(reagent.replace('self.',''))
                    else:
                        reag_l.append('{' + number_format % abs(reagents[reagent]) + '}' + reagent.replace('self.',''))
                elif reagents[reagent] == 0:
                    #reag_l.append(reagent.replace('self.',''))
                    print reagents
                    raw_input('WTF: please contact developers')

            if not reag_l:
                reag_l.append('$pool')
            if not reag_r:
                reag_r.append('$pool')

            substring = ''
            count = 0
            for x in reag_l:
                if count != 0:
                    substring += ' + '
                substring += x.replace(' ', '')
                count += 1
            prodstring = ''
            count = 0
            for x in reag_r:
                if count != 0:
                    prodstring += ' + '
                prodstring += x.replace(' ','')
                count += 1

            if NetworkDict[key]['Type'] == 'Rever':
                symbol = ' = '
            else:
                symbol = ' > '
            out_list.append('\t' + substring + symbol + prodstring + '\n')
            out_list.append('\t' + NetworkDict[key]['RateEq'].replace('self.','') + '\n\n')
        for x in out_list:
            self.psc_string += x

if __name__ == '__main__':
    #print SBML2psc(sbml_file='LV2.xml', dir_in='pscmodels').writePSCToString()
    print SBML2psc(sbml_file='LV2.xml', dir_in='pscmodels').writePSCToFile('tmp.psc')
